#pragma once

#include <utility>
#include <tbb/concurrent_queue.h>
#include "thread.h"
#include "basetypes.h"

class JobBase
{
public:
	// return false to re-enqueue
	virtual bool Execute() { return true; };

	virtual ~JobBase() { };
};

class ThreadPool
{
	template<typename F>
	class Job: public JobBase
	{
		F _func;
	public:
		Job(const F& func): _func(func) { }
		Job(F&& func): _func(std::move(func)) { }

		bool Execute() { return _func(); }

		~Job() { }
	};

	typedef tbb::concurrent_bounded_queue<JobBase*> ConcurrentQueue;
	typedef std::thread Thread;

	//size_t _maxThreadCount;
	size_t _threadCount;
	Thread*  _threads;
	ConcurrentQueue _queue;
	volatile bool isStopping;

public:
	ThreadPool(double maxRunningRate = 1.0)
	{
		isStopping = false;
		_threadCount = (size_t)(std::thread::hardware_concurrency() * maxRunningRate);
        if (_threadCount == 0)
        {
            _threadCount = 1;
        }
		_threads = (Thread*)malloc(sizeof(Thread) * _threadCount);

		for (size_t i = 0; i < _threadCount; ++i)
		{
			new(_threads + i) Thread([&] {
				while (!isStopping)
				{
					JobBase* pJob;
					_queue.pop(pJob);
					if (pJob->Execute())
						delete pJob;
					else
						_queue.push(pJob);
				}
			} );
		}
	}

	~ThreadPool() { StopAndWait(); }

	size_t ThreadCount() const { return _threadCount; }

	diff_t PendingJobCount() const { return _queue.size(); }

	bool IsStopping() const { return isStopping; }

	void Stop() { isStopping = true; }

	void StopAndWait() 
	{ 
		Stop();
		for (size_t i = 0; i < _threadCount; ++i) 
			_queue.push(new JobBase());
		for (size_t i = 0; i < _threadCount; ++i) 
			_threads[i].join();
	}

	// ThreadPool will delete Job after its Execute return true
	bool EnqueueJob(JobBase* pJob)
	{
		_queue.push(pJob);
		return true;
	}

	template<typename F>
	bool EnqueueJob(F&& f)
	{
		_queue.push(new Job<typename std::remove_reference<F>::type>(std::forward<F>(f)));
		return true;
	}
};
